package com.example.harmonet.harmtorch;

import ohos.global.resource.Resource;
import ohos.hiviewdfx.HiLog;
import ohos.hiviewdfx.HiLogLabel;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;

public class NN {
    ArrayList<Layer> _net;

    public NN(Layer... layers) {
        _net = new ArrayList<>();
        _net.addAll(Arrays.asList(layers));
    }

    public void readTorch(Resource f) {
        try {
            for (Layer layer : _net) {
                int param = layer.getParam();
                float[] params = new float[param];
                for(int i = 0; i < param; ++i) {
                    params[i] = readFloat(f);
                }
                layer.init(params);
            }
        } catch (IOException e) {
            HiLog.warn(new HiLogLabel(ohos.hiviewdfx.HiLog.LOG_APP, 0x00201, "test"), "read fail");
            // e.printStackTrace();
            System.exit(0);
        }
    }

    public Tensor inference(Tensor input) {
        Tensor out = input.view();
        try {
            for (Layer layer : _net) {
                out = layer.forward(out);
            }
        } catch (Exception e) {
            // e.printStackTrace();
        }
        return out;
    }

    static float readFloat(Resource f) throws IOException {
        byte[] bytes = new byte[4];
        f.read(bytes);
        String b = byteToArray(bytes);
        long ieee754Int = Long.parseLong(b, 16);
        return Float.intBitsToFloat((int) ieee754Int);
    }

    static String byteToArray(byte[] data){
        String result = "";
        for (int i = 0; i < data.length; i++) {
            result += Integer.toHexString((data[i] & 0xFF) | 0x100).toUpperCase().substring(1, 3);
        }
        return result;
    }
}
